package org.ricks.rpc.asm;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.lang.reflect.Method;
import java.net.ServerSocket;
import java.net.Socket;
import java.nio.ByteBuffer;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;

// 3. 服务端核心（虚拟线程支持）
public class RpcServer {
    private static final Map<String, Object> services = new ConcurrentHashMap<>();
    private static final Executor executor = Executors.newVirtualThreadPerTaskExecutor();

    public static void start(int port) throws IOException {
        try (ServerSocket server = new ServerSocket(port)) {
            System.out.println("Server started on port " + port);
            while (true) {
                Socket client = server.accept();
                executor.execute(() -> handleRequest(client));
            }
        }
    }

    private static void handleRequest(Socket client) {
        try (DataInputStream in = new DataInputStream(client.getInputStream());
             DataOutputStream out = new DataOutputStream(client.getOutputStream())) {

            // 协议头解析
            int magic = in.readInt();
            if (magic != 0xCAFEBABE) throw new IllegalArgumentException("Invalid protocol");

            // 请求体处理
            String serviceName = in.readUTF();
            String methodName = in.readUTF();
            int paramCount = in.readInt();

            Object[] args = new Object[paramCount];
            Class<?>[] paramTypes = new Class<?>[paramCount];
            for (int i = 0; i < paramCount; i++) {
                int paramLength = in.readInt();
                byte[] paramBytes = new byte[paramLength];
                in.readFully(paramBytes);
                args[i] = RpcCodec.DEFAULT_DECODER.apply(
                        ByteBuffer.wrap(paramBytes), Object.class);
                paramTypes[i] = args[i].getClass();
            }

            // 方法调用
            Object service = services.get(serviceName);
            Method method = service.getClass().getMethod(methodName, paramTypes);
            Object result = method.invoke(service, args);

            // 响应发送
            ByteBuffer resultBuffer = ByteBuffer.allocate(1024);
            RpcCodec.DEFAULT_ENCODER.accept(result, resultBuffer);
            out.writeInt(resultBuffer.position());
            out.write(resultBuffer.array(), 0, resultBuffer.position());

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static void registerService(Object service) {
        RpcService annotation = service.getClass().getAnnotation(RpcService.class);
        if (annotation != null) {
            services.put(annotation.serviceName(), service);
        }
    }
}